#include "track/post_processor.h"

#include <iostream>

namespace autodrive {
namespace perception {
Postprocessor::Postprocessor() {}

Postprocessor::~Postprocessor() {}

void Postprocessor::Init(const MapTrackParameter& settings,
                         const ProcessCategory& cat) noexcept {
  std::lock_guard<std::mutex> lock(mutex_);

  is_init_ = true;
  track_initialization_flag_ = false;
  settings_ = settings;
  pre_pose_.matrix() = Eigen::Matrix4d::Identity();
  cur_pose_.matrix() = Eigen::Matrix4d::Identity();
  fitting_ = std::make_unique<RectangleFitting>();
  matcher_ = std::make_unique<Mathcer>();
  dbscan_ = std::make_unique<DBScan<Eigen::Vector2f>>();
  if (cat == ProcessCategory::Turn) {
    dbscan_->SetEPS(settings.DBSCAN_eps_turn);
    dbscan_->SetMinNum(settings.DBSCAN_min_num_turn);
  } else if (cat == ProcessCategory::CrossWalk) {
    dbscan_->SetEPS(settings.DBSCAN_eps_crosswalk);
    dbscan_->SetMinNum(settings.DBSCAN_min_num_crosswalk);
  } else if (cat == ProcessCategory::No_Stop) {
    dbscan_->SetEPS(settings.DBSCAN_eps_nostop);
    dbscan_->SetMinNum(settings.DBSCAN_min_num_nostop);
  }
}

void Postprocessor::Process(
    autodrive::perception::Map& map,
    autodrive::perception::MapElementSet<RoadMarkingResult>& output_roadmarking,
    std::vector<Eigen::Vector2f>& ps, const Eigen::Affine3d& pos,
    const MapTrackParameter& settings) noexcept {
  std::string err_msg;
  if (!is_init_) {
    err_msg = "post processor module must do init  at first!";
    return;
  }

  if (ps.size() == 0) return;

  // A cluster with DBSCAN
  DBScan<Eigen::Vector2f> dbscan(1.0, 4);
  vector<vector<Eigen::Vector2f>> res =
      dbscan.GetClusters(ps);  // 第一个是噪点，其余的是聚类结果

  // vector<vector<Eigen::Vector2f>> res =
  //     dbscan_->GetClusters(ps);  // 第一个是噪点，其余的是聚类结果
  // B fit rotated rectangle
  vector<cv::Mat> mats;  // cv::Mat 是一个矩形框的四个角点
  mats.resize(res.size() - 1);
  vector<RotatedRect> rects;
  rects.resize(res.size() - 1);
  for (int i = 1; i < res.size(); ++i) {
    const vector<Eigen::Vector2f>& points = res.at(i);
    fitting_->Fit(points, rects.at(i - 1), mats.at(i - 1));
  }
  // C track
  cur_pose_ = pos;
  if (!track_initialization_flag_) {
    // 1 construct trackerobject
    std::vector<TrackedObjectPtr> tracked_objects;
    construct_tracked_objects(res, rects, mats, &tracked_objects);
    // 2 generate index
    std::vector<std::int32_t> objects_index;
    range(tracked_objects, objects_index);
    // 3 create new track
    create_new_tracks(tracked_objects, objects_index, object_tracks_);
    track_initialization_flag_ = true;
    pre_pose_ = cur_pose_;
  } else {
    Eigen::Matrix4f transformed_matrix =
        (cur_pose_.matrix().inverse() * pre_pose_.matrix())
            .cast<float>();                       // transformed matrix
    std::vector<Eigen::VectorXf> tracks_predict;  // track is null
    set_zero_predict(tracks_predict, object_tracks_.GetTracks().size());

    // 1 construct tracker objects
    std::vector<TrackedObjectPtr> tracked_objects;
    construct_tracked_objects(res, rects, mats, &tracked_objects);

    // 2 match
    std::vector<ObjectTrackPtr>& tracks = object_tracks_.GetTracks();
    std::vector<std::pair<int, int>> assignments;
    std::vector<int> unassigned_tracks;
    std::vector<int> unassigned_objects;
    matcher_->Match(&tracked_objects, tracks, tracks_predict, &assignments,
                    &unassigned_tracks, &unassigned_objects,
                    transformed_matrix);

    float time_diff = 0.0f;
    // 3 update with associated objects
    update_assigned_tracks(&tracked_objects, assignments, time_diff,
                           object_tracks_);

    // 4 update tracks without associated objects
    update_unassigned_tracks(tracks_predict, unassigned_tracks, time_diff,
                             object_tracks_);

    // 5 delete track
    delete_lost_tracks(object_tracks_);

    // 6 create new tracks for objects without associated tracks
    create_new_tracks(tracked_objects, unassigned_objects, object_tracks_);

    // 7 collect result
    collect_tracked_results(map, output_roadmarking, object_tracks_);

    // 8 set position
    pre_pose_ = cur_pose_;
  }
}

void Postprocessor::construct_tracked_objects(
    const vector<vector<Eigen::Vector2f>>& points,
    const vector<RotatedRect>& rects, const vector<cv::Mat>& mats,
    std::vector<TrackedObjectPtr>* tracked_objects) noexcept {
  int num_tracker = points.size();
  tracked_objects->clear();
  tracked_objects->resize(num_tracker - 1);
  for (int i = 0; i < num_tracker - 1; ++i) {
    TrackedObjectPtr tracker = std::make_shared<TrackedObject>();
    tracker->center =
        Eigen::Vector3f{rects.at(i).center.x, rects.at(i).center.y, 0.0f};
    tracker->length = rects.at(i).size.width;
    tracker->width = rects.at(i).size.height;
    tracker->height = 0.0f;
    tracker->yaw = rects.at(i).angle;  // deg
                                       // Add polygen
                                       // at<float>(0, 0), m.at<float>(0, 1)

    tracker->polygon.clear();
    for (int j = 0; j < 4; ++j) {
      if (tracker->polygon.full()) {
        break;
      }
      Eigen::Vector2f temp;
      temp.x() = mats.at(i).at<float>(j, 0);
      temp.y() = mats.at(i).at<float>(j, 1);

      tracker->polygon.add(temp);
    }

    const vector<Eigen::Vector2f>& p = points.at(i + 1);
    tracker->points.clear();
    for (int j = 0; j < p.size(); ++j) {
      if (tracker->points.full()) {
        break;
      }
      Eigen::Vector2f temp;
      temp.x() = p.at(j).x();
      temp.y() = p.at(j).y();
      tracker->points.add(temp);
    }
    (*tracked_objects)[i] = tracker;
  }
}

void Postprocessor::create_new_tracks(
    std::vector<TrackedObjectPtr> const& new_objects,
    std::vector<std::int32_t> const& unassigned_objects_index,
    ObjectTrackSet& object_tracks_set) {
  for (int i = 0; i < unassigned_objects_index.size(); ++i) {
    std::int32_t obj_id{unassigned_objects_index.at(i)};
    ObjectTrackPtr track(new ObjectTrack(new_objects[obj_id]));
    object_tracks_set.AddTrack(track);
  }
}

void Postprocessor::range(const std::vector<TrackedObjectPtr>& tracked_objects,
                          std::vector<std::int32_t>& objects_index) noexcept {
  objects_index.clear();
  std::int32_t num = tracked_objects.size();
  for (std::int32_t i = 0; i < num; ++i) {
    objects_index.push_back(i);
  }
}

void Postprocessor::update_assigned_tracks(
    std::vector<TrackedObjectPtr>* new_objects,
    std::vector<TrackObjectPair> const& assignments, double time_diff,
    ObjectTrackSet& object_tracks_set) {
  std::vector<ObjectTrackPtr>& tracks = object_tracks_set.GetTracks();
  for (size_t i = 0; i < assignments.size(); i++) {
    std::int32_t track_id = assignments[i].first;
    std::int32_t obj_id = assignments[i].second;
    tracks[track_id]->UpdateWithObject(&(*new_objects)[obj_id], time_diff,
                                       settings_);
  }
}

void Postprocessor::update_unassigned_tracks(
    std::vector<Eigen::VectorXf> const& tracks_predict,
    std::vector<std::int32_t> const& unassigned_tracks, double time_diff,
    ObjectTrackSet& object_tracks_set) {
  std::vector<ObjectTrackPtr>& tracks = object_tracks_set.GetTracks();
  for (size_t i = 0; i < unassigned_tracks.size(); ++i) {
    std::int32_t track_id = unassigned_tracks[i];
    tracks[track_id]->UpdateWithoutObject(tracks_predict[track_id], time_diff);
  }
}

void Postprocessor::delete_lost_tracks(ObjectTrackSet& object_tracks_set) {
  object_tracks_set.RemoveLostTracks();
}

void Postprocessor::collect_tracked_results(
    autodrive::perception::Map& perceptionResult,
    autodrive::perception::MapElementSet<RoadMarkingResult>& output_roadmarking,
    ObjectTrackSet& object_tracks_set) noexcept {
  const std::vector<ObjectTrackPtr>& tracks = object_tracks_set.GetTracks();
  for (std::uint32_t i{0}; i < tracks.size(); ++i) {
    TrackedObjectPtr result_obj{tracks[i]->current_object_};
    if (tracks[i]->consecutive_invisible_count_ >
        settings_.collect_track_consecutive_invisible_maximum) {
      continue;
    }

    if (tracks[i]->age_ < settings_.collect_track_age_minimum) {
      continue;
    }
    if (output_roadmarking.full()) {
      break;
    }
    // 将数值赋予输出结构体
    auto obs{output_roadmarking.end()};
    obs->id = tracks[i]->idx_;
    obs->center = result_obj->center;
    obs->length = result_obj->length;
    obs->width = result_obj->width;
    obs->height = result_obj->height;
    obs->yaw = result_obj->yaw;
    obs->score = result_obj->score;
    obs->polygon = result_obj->polygon;
    output_roadmarking.resize(output_roadmarking.size() + 1);
  }
}

void Postprocessor::set_zero_predict(
    std::vector<Eigen::VectorXf>& tracks_predict, int track_num) {
  tracks_predict.resize(track_num);
  for (int i = 0; i < track_num; ++i) {
    Eigen::VectorXf& vec = tracks_predict.at(i);
    vec.resize(7);
    vec(0) = 0.0f;
    vec(1) = 0.0f;
    vec(2) = 0.0f;
    vec(3) = 0.0f;
    vec(4) = 0.0f;
    vec(5) = 0.0f;
    vec(6) = 0.0f;
  }
}

}  // namespace perception
}  // namespace autodrive
